"""Define the configurable parameters for the agent."""

from __future__ import annotations

from dataclasses import dataclass, field, fields
from typing import Annotated, Optional

from langchain_core.runnables import RunnableConfig, ensure_config

from react_agent import prompts
import os


def get_museumai_root() -> str:
    """Get museumai repo root path."""
    root_path = os.path.join(os.path.dirname(__file__), '..',
                             '..', '..').replace("\\", "/")
    if os.path.exists(root_path):
        print("Museumai root found: " + root_path)
        return root_path
    else:
        raise ValueError(
            f"Museumai root path not found.")


MUSEUMAI_ROOT = get_museumai_root()


@dataclass(kw_only=True)
class Configuration:
    """The configuration for the agent."""

    text_qa_prompt: str = field(
        default=prompts.TEXT_QA_PROMPT,
        metadata={
            "description": "The system prompt to use for the agent's interactions. "
            "This prompt sets the context and behavior for the agent."
        },
    )

    image_qa_prompt: str = field(
        default=prompts.IMAGE_QA_PROMPT,
        metadata={
            "description": "The image qa prompt to use for the agent's interactions. "
        },
    )

    model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = field(
        default="llama3.1:8b",
        metadata={
            "description": "The name of the language model to use for the agent's main interactions. "
            "Should be ollama model name"
        },
    )

    model_max_predict: int = field(
        default=512,
        metadata={
            "description": "The maximum output tokens number of model predict."
        },
    )

    model_context: int = field(
        default=4096,
        metadata={
            "description": "The maximum context limit for model run."
        },
    )

    max_search_results: int = field(
        default=10,
        metadata={
            "description": "The maximum number of search results to return for each search query."
        },
    )

    museumai_root: str = field(
        default=MUSEUMAI_ROOT,
        metadata={
            "description": "The root path of the museumai repository."
        },
    )

    text_db_path: str = field(
        default=MUSEUMAI_ROOT + "/data/dunhuang_text_db",
        metadata={
            "description": "The path to the text knowledge base for the agent to retrieve text from."
        },
    )

    text_top_k: int = field(
        default=5,
        metadata={
            "description": "The number of top-k text retrieved from the text knowledge base."
        },
    )

    img_db_path: str = field(
        default=MUSEUMAI_ROOT + "/data/dunhuang_db",
        metadata={
            "description": "The path to the image knowledge base for the agent to retrieve image from."
        },
    )

    img_top_k: int = field(
        default=1,
        metadata={
            "description": "The number of top-k image retrieved from the image knowledge base."
        },
    )

    @classmethod
    def from_runnable_config(
        cls, config: Optional[RunnableConfig] = None
    ) -> Configuration:
        """Create a Configuration instance from a RunnableConfig object."""
        config = ensure_config(config)
        configurable = config.get("configurable") or {}
        _fields = {f.name for f in fields(cls) if f.init}
        return cls(**{k: v for k, v in configurable.items() if k in _fields})
